"""
《邢不行-2023新版|Python股票量化投资课程》
author: 邢不行
微信: xbx9585

根据选股数据，进行选股
"""

from Config import *
from Evaluate import *
from Filter import *
from Functions import *
import warnings
import time

warnings.filterwarnings('ignore')

pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
pd.set_option('display.max_rows', 5000)  # 最多显示数据的行数

print('策略名称:', strategy_name)
print('周期:', period_type)

now = time.localtime()
nowt = time.strftime("%Y-%m-%d %H:%M:%S", now)
filename_day = time.strftime("%Y-%m-%d", now)

# ===导入数据
# 从pickle文件中读取整理好的所有股票数据
df = pd.read_pickle(root_path + '/data/output/选股策略/all_stock_data_%s.pkl' % period_type)
#注释可下周选股
df.dropna(subset=['下周期每天涨跌幅'], inplace=True)
# 导入指数数据
index_data = import_index_data(root_path + '/data/index_data/sh000300.csv', back_trader_start=date_start,
                               back_trader_end=date_end)

# 创造空的事件周期表，用于填充不选股的周期
empty_df = create_empty_data(index_data, period_type)

# ===删除新股
df = df[df['上市至今交易天数'] > 250]
# 删除创业板
sz30 = df[df['股票代码'].str.contains('sz30')]
df = df.drop(sz30.index)

# ===删除下个交易日不交易、开盘涨停的股票，因为这些股票在下个交易日开盘时不能买入。
df = df[df['下日_是否交易'] == 1]
df = df[df['下日_开盘涨停'] == False]
df = df[df['下日_是否ST'] == False]
df = df[df['下日_是否退市'] == False]
# df = df[df['交易日期'] == '2023/4/7']
# print(df[['交易日期','股票代码','股票名称','成交额','zs_signal']].tail(100))
# exit()

# 选股
df = filter_and_rank(df, stg_name=strategy_name)
df['zs_signal'] = df['zs_signal'].fillna(0)
# print(df[['交易日期','股票代码','股票名称','成交额','zs_signal']].tail(10))
# exit()
# === 以下数据更新方式二选一
# === 1、每次执行会生成以日期命名的csv文件
print("更新时间：", nowt)
path = root_path+'/data/output/选股策略/' + filename_day + strategy_name +'.csv'
pd.DataFrame(columns=['更新时间：' + nowt]).to_csv(path, index=False, encoding='gbk')
df.to_csv(path, index=False, mode='a', encoding='gbk')

select_data = df
# print(df[['交易日期','股票代码','股票名称','zs_signal']].tail(10))
# exit()

# ===选股
# ===按照开盘买入的方式，修正选中股票在下周期每天的涨跌幅。
# 即将下周期每天的涨跌幅中第一天的涨跌幅，改成由开盘买入的涨跌幅
df['下日_开盘买入涨跌幅'] = df['下日_开盘买入涨跌幅'].apply(lambda x: [x])
df['下周期每天涨跌幅'] = df['下周期每天涨跌幅'].apply(lambda x: x[1:])
df['下周期每天涨跌幅'] = df['下日_开盘买入涨跌幅'] + df['下周期每天涨跌幅']

# ===整理选中股票数据
# 挑选出选中股票
df['股票代码'] += ' '
df['股票名称'] += ' '
group = df.groupby('交易日期')
select_stock = pd.DataFrame()
# print(group.tail(10))
# exit()
select_stock['股票数量'] = group['股票名称'].size()
select_stock['买入股票代码'] = group['股票代码'].sum()
select_stock['买入股票名称'] = group['股票名称'].sum()
select_stock['是否开仓'] = group['zs_signal'].sum()
# print(select_stock)
# exit()

# 计算下周期每天的资金曲线
select_stock['选股下周期每天资金曲线'] = group['下周期每天涨跌幅'].apply(lambda x: np.cumprod(np.array(list(x)) + 1, axis=1).mean(axis=0))
# x = df.iloc[:3]['下周期每天涨跌幅']
# print()
# print(x)
# print(list(x))  # 将x变成list
# print(np.array(list(x)))  # 矩阵化
# print(np.array(list(x)) + 1)  # 矩阵中所有元素+1
# print(np.cumprod(np.array(list(x)) + 1, axis=1))  # 连乘，计算资金曲线
# print(np.cumprod(np.array(list(x)) + 1, axis=1).mean(axis=0))  # 连乘，计算资金曲线

# 扣除买入手续费
select_stock['选股下周期每天资金曲线'] = select_stock['选股下周期每天资金曲线'] * (1 - c_rate)  # 计算有不精准的地方
# 扣除卖出手续费、印花税。最后一天的资金曲线值，扣除印花税、手续费
select_stock['选股下周期每天资金曲线'] = select_stock['选股下周期每天资金曲线'].apply(
    lambda x: list(x[:-1]) + [x[-1] * (1 - c_rate - t_rate)])

# 计算下周期整体涨跌幅
select_stock['选股下周期涨跌幅'] = select_stock['选股下周期每天资金曲线'].apply(lambda x: x[-1] - 1)
# 计算下周期每天的涨跌幅
select_stock['选股下周期每天涨跌幅'] = select_stock['选股下周期每天资金曲线'].apply(
    lambda x: list(pd.DataFrame([1] + x).pct_change()[0].iloc[1:]))
del select_stock['选股下周期每天资金曲线']
# print(select_stock.tail(10))
# exit()
# 将选股结果更新到empty_df上

empty_df.update(select_stock)
select_stock = empty_df
# print(select_stock)
# exit()


# 从第一次买入股票的时候开始回测策略
select_stock = select_stock[select_stock['股票数量'].expanding().sum() > 0]

# 计算整体资金曲线
select_stock.reset_index(inplace=True)
select_stock['资金曲线'] = (select_stock['选股下周期涨跌幅'] + 1).cumprod()
print(select_stock)
# exit()

# 2024-1-20 存储每周涨幅pkl用于择时回测
select_stock.to_pickle(root_path + '/data/equity_curve/每周涨跌幅_%s_%s_%s.pkl' % (strategy_name, select_stock_num, period_type))

# ===计算选中股票每天的资金曲线
# 计算每日资金曲线
equity = pd.merge(left=index_data, right=select_stock[['交易日期', '买入股票代码','是否开仓','成交额']], on=['交易日期'],
                  how='left', sort=True)  # 将选股结果和大盘指数合并


equity['持有股票代码'] = equity['买入股票代码'].shift()
equity['持有股票代码'].fillna(method='ffill', inplace=True)
equity.dropna(subset=['持有股票代码'], inplace=True)
del equity['买入股票代码']
equity['持有股票代码是否开仓'] = equity['是否开仓'].shift()
equity['持有股票代码是否开仓'].fillna(method='ffill', inplace=True)
# equity.dropna(subset=['持有股票代码是否开仓'], inplace=True)
del equity['是否开仓']
equity['涨跌幅'] = select_stock['选股下周期每天涨跌幅'].sum()
equity['equity_curve'] = (equity['涨跌幅'] + 1).cumprod()
equity['benchmark'] = (equity['指数涨跌幅'] + 1).cumprod()

# # 周均线择时
# equity['短周均']=equity['equity_curve'].rolling(3, min_periods=1).mean()
# equity['长周均']=equity['equity_curve'].rolling(19, min_periods=1).mean()
# condition1 = equity['短周均'] > equity['长周均']  # 短期均线 > 长期均线
# condition2 = equity['短周均'].shift(1) <= equity['长周均'].shift(1)  # 上一周期的短期均线 <= 长期均线
# equity.loc[condition1, 'signal'] = 1  # 将产生做多信号的那根K线的signal设置为1，1代表做多

# # ===找出做多平仓信号
# condition1 = equity['短周均'] < equity['长周均']  # 短期均线 < 长期均线
# condition2 = equity['短周均'].shift(1) >= equity['长周均'].shift(1)  # 上一周期的短期均线 >= 长期均线
# equity.loc[condition1, 'signal'] = 0  # 将产生平仓信号当天的signal设置为0，0代表平仓

# 4周均线择时
equity['4周均']=equity['equity_curve'].rolling(4, min_periods=1).mean()
condition1 = equity['equity_curve'] >= equity['4周均']  # 短期均线 > 长期均线
equity.loc[condition1, 'signal'] = 1  # 将产生做多信号的那根K线的signal设置为1，1代表做多

# ===找出做多平仓信号
condition1 = equity['equity_curve'] < equity['4周均']  # 短期均线 < 长期均线
equity.loc[condition1, 'signal'] = 0  # 将产生平仓信号当天的signal设置为0，0代表平仓

# 计算持仓
# equity['signal'].fillna(method='ffill', inplace=True)
# equity['signal'].fillna(value=0, inplace=True)
# equity['pos'] = equity['signal'].shift(1)
# equity2= equity[equity['pos'] == 1]
# equity['择时后净值']=(equity2['涨跌幅']+1).cumprod()
# equity.fillna(method='ffill', inplace=True)
print(equity.tail(20))


# equity2['equity_curve'] =  (equity2['涨跌幅']+1).cumprod()
# equity= equity[equity['信号']]
# equity['择时后净值']=(equity['涨跌幅']+1).cumprod()

# path = root_path+'/data/output/选股策略/' + filename_day + strategy_name +'净值.csv'
# pd.DataFrame(columns=['更新时间：' + nowt]).to_csv(path, index=False, encoding='gbk')
# equity.to_csv(path, index=False, mode='w', encoding='gbk')
# path2 = root_path+'/data/output/选股策略/' + filename_day + strategy_name +'择时净值.csv'
# pd.DataFrame(columns=['更新时间：' + nowt]).to_csv(path, index=False, encoding='gbk')
# equity2.to_csv(path2, index=False, mode='w', encoding='gbk')
# print(equity2)
# exit()


def market_timing_para_instant(select_stock_num, select_data, equity, tittle='是否开仓'):
    """
    选股时，根据导入参指数数择时参数，如果参数为正向则买入股票，如果参数为反向则空仓

    :param tittle: select_data的择时信号，如20日涨跌幅，如MA20等信号
    :param time: select_data的交易日期
    :param select_data:选股文件，选出交易股票参数的文件
    :param equity: 已经处理好的资金曲线，继续后续处理
    :param select_stock_num:选股数量
    :return:
    """

    equity = equity.copy()

    # print(select_data[['交易日期','股票代码','zs_signal']].tail(10))
    # exit()
    # 计算择时信号，当持仓50%股票发出空仓信号，则空仓
    equity['持有股票代码_择时'] = equity['持有股票代码'].copy()
    # equity['是否开仓'] = select_data[tittle].sum()
    # equity['是否开仓'] = equity['是否开仓'].fillna(method='bfill')
  
    equity['持有股票代码是否开仓'] = equity['持有股票代码是否开仓'] / select_stock_num
  
    equity.loc[equity['持有股票代码是否开仓'] < 0.5, '持有股票代码_择时'] = '空仓'
  
    # equity['持有股票代码_择时'] = equity['持有股票代码_择时'].shift(1)
    # print(equity.tail(20))
    # exit()
    equity.loc[equity['持有股票代码_择时'] == '空仓', '持有股票代码'] = '空仓'
    del equity['持有股票代码_择时']

    equity['涨跌幅_择时'] = equity['涨跌幅'].copy()
    equity.loc[equity['持有股票代码'] == '空仓', '涨跌幅_择时'] = 0
    # 计算择时涨跌幅
    equity['equity_curve_' + tittle + '择时'] = (equity['涨跌幅_择时'] + 1).cumprod()

    return equity



# equity2 = market_timing_para_instant(select_stock_num,select_data,equity,'zs_signal')
# equity2 = RSIV_signal(equity)
# equity['择时后净值'] = equity2['equity_curve_' + 'zs_signal' + '择时']

# ===计算策略评价指标
rtn, year_return, month_return = strategy_evaluate(equity, select_stock)
# rtn, year_return, month_return = strategy_evaluate(equity2, select_stock)
print(rtn)

# ===画图
equity = equity.reset_index()
# draw_equity_curve_mat(equity, data_dict={'策略表现': 'equity_curve', '基准涨跌幅': 'benchmark'}, date_col='交易日期')
draw_equity_curve_mat(equity, data_dict={'策略表现': 'equity_curve', '基准涨跌幅': 'benchmark'}, date_col='交易日期')
# 如果上面的函数不能画图，就用下面的画图
# draw_equity_curve_plotly(equity, data_dict={'策略涨跌幅': 'equity_curve', '基准涨跌幅': 'benchmark'}, date_col='交易日期')
